import ast
import json
from typing import Tuple, Optional, Dict

# Load model configurations from JSON
with open('media_configs_model.json', 'r') as f:
    model_configs = json.load(f)

# Load template configurations from JSON
with open('media_configs_template.json', 'r') as f:
    template_configs = json.load(f)

HUB_R2_STORE_PREFIX = "video/hub_video"

# Default watermark and thumbnail settings
SETTING_WATERMARK_THUMBNAIL_INFO = {
    "thumbnail_time_seconds": 0,
    "watermark_position": "top-right",
    "watermark_offset_x": 20,
    "watermark_offset_y": 15
}

# AIHubVideoTasks
# video_media_thumbmark
# version 2
TASK_QUEUE_AIGCVIDEO_NAME = "AIHubVideoTasks2"
TASK_QUEUE_THUMBMART_NAME = "video_media_thumbmark2"

# TASK_QUEUE_AIGCVIDEO_NAME = "AIHubVideoTasks"
# TASK_QUEUE_THUMBMART_NAME = "video_media_thumbmark"

# NSFW 定义各内容类别的阈值
RESOLUTION_CONTENT_REVIEW = {
    "sexual": 0.85,
    # "sexual/minors": 1,
    # "Harassment": 1,
    # "harassment/threatening": 1,
    # "Hate": 1,
    # "hate/threatening": 1,
    # "illicit": 1,
    # "illicit/violent": 1,
    # "self-harm":1,
    # "self-harm/intent": 1,
    # "self-harm/instructions": 1,
    "violence": 0.83,
    # "violence/graphic": 1,
}

# Valid aspect ratios and resolutions
VALID_ASPECT_RATIOS = ['16:9', '4:3', '1:1', '3:4', '9:16', '21:9']
VALID_RESOLUTIONS = ['360p', '480p', '540p', '720p', '768p', '1080p', '2k', '4k']

# Resolution reference data
RESOLUTION_REFERENCE = {
    '360p': {'4:3': (480, 360), '15:9': (600, 360), "max_side": 480},
    '480p': {'16:9': (854, 480), '4:3': (736, 544), '1:1': (640, 640), '3:4': (544, 736), '9:16': (480, 864), '21:9': (960, 416), "max_side": 854},
    '540p': {'16:9': (960, 540), "max_side": 960},
    '720p': {'16:9': (1248, 704), '4:3': (1120, 832), '1:1': (960, 960), '3:4': (832, 1120), '9:16': (704, 1248), '21:9': (1504, 640), "max_side": 1280},
    '768p': {'16:9': (1248, 704), '4:3': (1120, 832), '1:1': (960, 960), '3:4': (832, 1120), '9:16': (704, 1248), '21:9': (1504, 640), "max_side": 1366},
    '1080p': {'16:9': (1920, 1088), '4:3': (1664, 1248), '1:1': (1440, 1440), '3:4': (1248, 1664), '9:16': (1088, 1920), '21:9': (2176, 928), "max_side": 1920},
    '2k': {'max_side': 2560},
    '4k': {'max_side': 3840}
}

RESOLUTION_MAP = {
    '360p': {'max_side': 480},
    '480p': {'max_side': 854},
    '540p': {'max_side': 960},
    '720p': {'max_side': 1280},
    '768p': {'max_side': 1366},
    '1080p': {'max_side': 1920},
    '2k': {'max_side': 2560},
    '4k': {'max_side': 3840}
}

# Map resolution to base value
RESOLUTION_BASE_VALUE = {
    '360p': 360,
    '480p': 480,
    '540p': 540,
    '720p': 720,
    '768p': 768,
    '1080p': 1080
}


def get_poll_config(model: str) -> Dict[str, int]:
    """获取指定模型的结果轮询时间配置。

    Args:
        model: 模型名称。

    Returns:
        包含超时时间和轮询间隔的字典。

    Raises:
        ValueError: 如果模型不支持。
    """
    model_poll_configs = {
        'doubao': {'timeout': 600, 'poll_interval': 5},
        'google': {'timeout': 1800, 'poll_interval': 20}
    }
    return model_poll_configs[model]


def calculate_template_ratio_resolution(module_type: str, duration: str, target_quality: str = None, custom_width: Optional[int] = None, custom_height: Optional[int] = None) -> Dict[str, any]:
    """
    根据module_type和duration从模板配置中获取multiplier值和目标分辨率

    :param module_type: 模块类型
    :param duration: 时长
    :param target_quality: 目标分辨率档位，如'540p', '720p'等
    :param custom_width: 自定义宽度
    :param custom_height: 自定义高度
    :return: 包含multiplier和target_resolution的字典
    :raises ValueError: 当找不到对应配置时
    """

    # 只有当duration不以"s"结尾时才添加"s"
    if not isinstance(duration, str) or not duration.endswith("s"):
        duration = f"{duration}s"
    
    print(f"模板计算超分参数: module_type={module_type}, duration={duration},target_quality={target_quality}, custom_width={custom_width}, custom_height={custom_height}")
    "模板计算超分参数: module_type=22, duration=5s,target_quality=480p, custom_width=1024, custom_height=1280"

    # 检查module_type是否存在于模板配置中
    if str(module_type) not in template_configs:
        raise ValueError(f"不支持的module_type: {module_type}")

    module_config = template_configs[str(module_type)]

    # 检查duration是否存在于配置中
    if "durations" not in module_config or duration not in module_config["durations"]:
        raise ValueError(f"不支持的duration: {duration}")

    duration_config = module_config["durations"][duration]

    # 如果未指定target_quality，则使用配置中的第一个分辨率
    if target_quality is None:
        if "resolutions" in duration_config and duration_config["resolutions"]:
            target_quality = list(duration_config["resolutions"].keys())[0]
        else:
            raise ValueError(f"在duration {duration} 中未找到任何分辨率配置")

    # 检查resolutions是否存在
    if "resolutions" not in duration_config or target_quality not in duration_config["resolutions"]:
        raise ValueError(f"不支持的target_quality: {target_quality}")

    resolution_config = duration_config["resolutions"][target_quality]

    # 获取multiplier值
    multiplier = resolution_config.get("multiplier", 1)
    base_resolution = resolution_config['base_resolution']
    max_side = resolution_config['max_side']
    base_value = RESOLUTION_BASE_VALUE.get(base_resolution, 720)  # Default to 720 if not found

    # 计算目标分辨率
    target_width = custom_width * multiplier if custom_width else None
    target_height = custom_height * multiplier if custom_height else None

    # Calculate super resolution
    if custom_width and custom_height:
        super_width = custom_width * multiplier
        super_height = custom_height * multiplier
    else:
        raise ValueError("自定义尺寸和目标分辨率不能同时为空")

    # 如果有自定义尺寸，根据实际宽高比和RESOLUTION_MAP中的max_side重新计算target_resolution
    if super_width and super_height:
        target_width, target_height = calculate_max_side_resolution(super_width, super_height, max_side)

    return {
        'multiplier': multiplier,
        'target_resolution': (target_width, target_height),
        'base_resolution': base_value,
        'max_side': resolution_config.get("max_side")
    }


def calculate_aspect_ratio_resolution(aspect_ratio: str, max_side: int) -> Tuple[int, int]:
    """根据宽高比和最大边计算分辨率。

    Args:
        aspect_ratio: 宽高比，例如 '16:9'。
        max_side: 最长边的像素值。

    Returns:
        宽度和高度的元组 (width, height)。

    Raises:
        ValueError: 如果宽高比无效。
    """
    if aspect_ratio not in VALID_ASPECT_RATIOS:
        raise ValueError(f"不支持的宽高比: {aspect_ratio}，支持的宽高比为: {', '.join(VALID_ASPECT_RATIOS)}")

    width_ratio, height_ratio = map(int, aspect_ratio.split(':'))
    if width_ratio > height_ratio:
        width = max_side
        height = int(max_side * height_ratio / width_ratio)
    else:
        height = max_side
        width = int(max_side * width_ratio / height_ratio)
    return (width, height)


def calculate_max_side_resolution(width: int, height: int, max_side: int) -> Tuple[int, int]:
    """根据实际宽高和最大边计算保持宽高比的分辨率。

    Args:
        width: 实际宽度
        height: 实际高度
        max_side: 最大边长限制

    Returns:
        保持宽高比且最大边不超过max_side的宽高元组 (width, height)
    """
    # 计算宽高比
    aspect_ratio = width / height

    if width > height:
        new_width = min(width, max_side)
        new_height = int(new_width / aspect_ratio)
        # 确保高度不超过max_side
        if new_height > max_side:
            new_height = max_side
            new_width = int(new_height * aspect_ratio)
    else:
        new_height = min(height, max_side)
        new_width = int(new_height * aspect_ratio)
        # 确保宽度不超过max_side
        if new_width > max_side:
            new_width = max_side
            new_height = int(new_width / aspect_ratio)

    # 确保宽度和高度都是偶数
    if new_width % 2 != 0:
        new_width -= 1
    if new_height % 2 != 0:
        new_height -= 1

    return (new_width, new_height)


def calculate_scaled_dimensions(original_width: int, original_height: int, quality: str, aspect_ratio: str, multiplier: float = 1) -> Tuple[int, int]:
    """计算图片缩放后的宽高参数。

    Args:
        original_width: 原始图片宽度。
        original_height: 原始图片高度。
        quality: 分辨率档位，例如 '480p'。
        aspect_ratio: 宽高比，例如 '16:9'。
        multiplier: 放大倍数，默认为 1。

    Returns:
        缩放后的宽度和高度元组 (width, height)。

    Raises:
        ValueError: 如果档位或宽高比不支持。
    """
    if quality not in RESOLUTION_MAP:
        raise ValueError(f"不支持的档位: {quality}，支持的档位为: {', '.join(RESOLUTION_MAP.keys())}")

    if aspect_ratio == "adaptive":
        aspect_ratio = "4:3"

    if aspect_ratio not in VALID_ASPECT_RATIOS:
        raise ValueError(f"不支持的宽高比: {aspect_ratio}，支持的宽高比为: {', '.join(VALID_ASPECT_RATIOS)}")

    scaled_width = int(original_width * multiplier)
    scaled_height = int(original_height * multiplier)
    max_side = RESOLUTION_MAP[quality]['max_side']
    target_width, target_height = calculate_aspect_ratio_resolution(aspect_ratio, max_side)

    if scaled_width > target_width or scaled_height > target_height:
        ratio = min(target_width / scaled_width, target_height / scaled_height)
        scaled_width = round(scaled_width * ratio)
        scaled_height = round(scaled_height * ratio)

    return (scaled_width, scaled_height)


def calculate_closest_aspect_ratio(custom_width: Optional[int], custom_height: Optional[int]) -> str:
    """根据自定义宽高计算最接近的预定义宽高比。

    Args:
        custom_width: 自定义宽度。
        custom_height: 自定义高度。

    Returns:
        最接近的预定义宽高比字符串，例如 '16:9'。
    """
    if not custom_width or not custom_height:
        return "16:9"

    actual_ratio = custom_width / custom_height
    closest_ratio = VALID_ASPECT_RATIOS[0]
    min_difference = float('inf')

    for ratio in VALID_ASPECT_RATIOS:
        width_ratio, height_ratio = map(int, ratio.split(':'))
        target_ratio = width_ratio / height_ratio
        difference = abs(actual_ratio - target_ratio)
        if difference < min_difference:
            min_difference = difference
            closest_ratio = ratio

    return closest_ratio


def calculate_super_resolution(
        model: str,
        target_quality: str,
        aspect_ratio: str,
        duration: str,
        mode: str = 'std',
        custom_width: Optional[int] = None,
        custom_height: Optional[int] = None,
        module_type: Optional[int] = None
) -> Dict[str, any]:
    """根据目标档位、宽高比、时长和模式计算超分参数及结果。

    Args:
        model: 模型名称，例如 'doubao-seedance-1-0-pro'。
        target_quality: 目标分辨率档位，例如 '480p'。
        aspect_ratio: 宽高比，例如 '16:9' 或 'adaptive'。
        duration: 时长，例如 '5s' 或 '10s'。
        mode: 模式，'std' 或 'pro'，默认为 'std'。
        custom_width: 可选的自定义宽度。
        custom_height: 可选的自定义高度。
        module_type: 可选的模块类型，用于特定模型映射。

    Returns:
        包含超分信息的字典，包括基础档位、放大倍数、超分后尺寸、目标尺寸。

    Raises:
        ValueError: 如果模型、档位、宽高比、时长或模式不支持。
    """
    duration = f"{duration}s"
    print(f"计算超分参数: model={model}, quality={target_quality}, aspect_ratio={aspect_ratio}, duration={duration}, mode={mode}, custom_width={custom_width}, custom_height={custom_height}")

    # Handle module_type for model mapping
    if module_type == 22:
        model = "v3.5-v4-v4.5"

    # Validate model
    if model not in model_configs:
        raise ValueError(f"不支持的模型: {model}，支持的模型为: {', '.join(model_configs.keys())}")

    # Validate mode
    if mode not in model_configs[model]['modes']:
        raise ValueError(f"模型 {model} 不支持的模式: {mode}，支持的模式为: {', '.join(model_configs[model]['modes'].keys())}")

    # Validate duration
    if duration not in model_configs[model]['modes'][mode]['durations']:
        raise ValueError(f"模型 {model} 模式 {mode} 不支持的时长: {duration}，支持的时长为: {', '.join(model_configs[model]['modes'][mode]['durations'].keys())}")
    target_quality = target_quality.lower()
    # Validate resolution
    quality_dependencies = model_configs[model]['modes'][mode]['durations'][duration]['resolutions']
    if target_quality not in quality_dependencies:
        raise ValueError(f"模型 {model} 模式 {mode} 时长 {duration} 不支持的目标档位: {target_quality}，支持的档位为: {', '.join(quality_dependencies.keys())}")

    # Handle adaptive aspect ratio
    if not aspect_ratio or aspect_ratio == "adaptive":
        aspect_ratio = calculate_closest_aspect_ratio(custom_width, custom_height)
        print(f"自适应宽高比: {aspect_ratio}")

    # Get configuration
    config = quality_dependencies[target_quality]
    base_resolution = config['base_resolution']
    multiplier = config['multiplier']
    max_side = config['max_side']

    # Calculate base value from base_resolution
    base_value = RESOLUTION_BASE_VALUE.get(base_resolution, 720)  # Default to 720 if not found

    # Calculate target resolution
    if target_quality in RESOLUTION_REFERENCE and aspect_ratio in RESOLUTION_REFERENCE[target_quality]:
        target_width, target_height = RESOLUTION_REFERENCE[target_quality][aspect_ratio]
    else:
        target_width, target_height = calculate_aspect_ratio_resolution(aspect_ratio, max_side)

    # Calculate super resolution
    if custom_width and custom_height:
        super_width = custom_width * multiplier
        super_height = custom_height * multiplier
    else:
        base_width, base_height = calculate_aspect_ratio_resolution(aspect_ratio, base_value)
        super_width = base_width * multiplier
        super_height = base_height * multiplier

        if super_width > target_width or super_height > target_height:
            ratio = min(target_width / super_width, target_height / super_height)
            super_width = max(1, round(super_width * ratio))
            super_height = max(1, round(super_height * ratio))

    # 如果有自定义尺寸，根据实际宽高比和RESOLUTION_MAP中的max_side重新计算target_resolution
    if super_width and super_height:
        target_width, target_height = calculate_max_side_resolution(super_width, super_height, max_side)

    return {
        'target_quality': target_quality,
        'base_resolution': base_resolution,
        'base_value': base_value,
        'multiplier': multiplier,
        'max_side': max_side,
        'custom_size': (custom_width, custom_height) if custom_width and custom_height else None,
        'super_resolution': (super_width, super_height),
        'target_resolution': (target_width, target_height)
    }
